"""
@Time ： 2025/9/18 10:38
@Auth ： zh
@File ：db_utils.py
"""
"""
@Time ： 2025/9/18 10:30
@Auth ： zh
@File ：db_utils.py
"""
import pymysql
from pymysql import OperationalError, MySQLError
from pymysql.cursors import DictCursor


class MySQLConnector:
    """MySQL数据库连接工具类"""

    def __init__(self, host_name, user_name, user_password, db_name=None, port=3306):
        self.host = host_name
        self.user = user_name
        self.password = user_password
        self.db_name = db_name
        self.port = port
        self.connection = None

    def connect(self):
        """建立数据库连接"""
        try:
            connection_params = {
                'host': self.host,
                'user': self.user,
                'password': self.password,
                'port': self.port,
                'charset': 'utf8mb4'
            }
            if self.db_name:
                connection_params['db'] = self.db_name

            self.connection = pymysql.connect(**connection_params)

            if self.connection.open:
                db_info = self.connection.get_server_info()
                print(f"成功连接到MySQL服务器，版本: {db_info}")

                if self.db_name:
                    with self.connection.cursor() as cursor:
                        cursor.execute("select database();")
                        record = cursor.fetchone()
                        print(f"当前连接的数据库: {record[0]}")
            return True

        except OperationalError as e:
            print(f"连接错误: {e}")
            self.connection = None
            return False
        except Exception as e:
            print(f"发生错误: {e}")
            self.connection = None
            return False

    def close(self):
        """关闭数据库连接"""
        if self.connection and self.connection.open:
            self.connection.close()
            print("数据库连接已关闭")
            self.connection = None

    def execute_query(self, query, params=None, use_dict_cursor=True):
        """执行查询语句"""
        if not self.connection or not self.connection.open:
            print("数据库连接未建立或已关闭")
            return []

        try:
            cursor_class = DictCursor if use_dict_cursor else None
            with self.connection.cursor(cursor_class) as cursor:
                cursor.execute(query, params or ())
                results = cursor.fetchall()
                print(f"查询成功，共找到 {len(results)} 条记录")
                return results

        except MySQLError as e:
            print(f"查询错误: {e}")
            return []

    def execute(self, query, params=None):
        """执行增删改等非查询语句"""
        if not self.connection or not self.connection.open:
            print("数据库连接未建立或已关闭")
            return False

        try:
            with self.connection.cursor() as cursor:
                affected_rows = cursor.execute(query, params or ())
                self.connection.commit()  # 提交事务
                print(f"执行成功，影响行数: {affected_rows}")
                return True
        except MySQLError as e:
            self.connection.rollback()  # 出错时回滚
            print(f"执行错误: {e}")
            return False

    def __enter__(self):
        """支持上下文管理器协议"""
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """自动关闭连接"""
        self.close()